package org.example.slave.executor.mysql.curd;

import lombok.SneakyThrows;
import org.example.slave.channel.MyChannel;
import org.example.slave.connector.mysql.MysqlConnector;
import org.example.slave.packet.*;
import org.example.slave.packet.execute.FieldPacket;
import org.example.slave.packet.execute.ResultSetPacket;
import org.example.slave.util.MyChannelUtil;
import org.example.slave.util.MysqlCodecUtil;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

/**
 * @author dingrui
 * @since 2022/2/16
 * @description
 */
public class MysqlQueryExecutor {

    private MysqlConnector connector;

    public MysqlQueryExecutor(MysqlConnector connector) {
        this.connector = connector;
    }

    protected byte[] readNextPacket() throws IOException {
        MyChannel ch = this.connector.getMyChannel();
        HeaderPacket h = MysqlCodecUtil.readHeader(ch, 4);
        return MyChannelUtil.readBytes(ch, h.getPacketBodyLength());
    }

    @SneakyThrows
    public ResultSetPacket query(String command) {
        QueryCommandPacket cmd = new QueryCommandPacket();
        cmd.setQueryString(command);
        byte[] bodyBytes = cmd.toBytes();
        MyChannel channel = this.connector.getMyChannel();
        MyChannelUtil.write(channel, bodyBytes, 0);
        byte[] body = this.readNextPacket();

        if (body[0] < 0) {
            ErrPacket errPacket = new ErrPacket(body);
            throw new RuntimeException("err: " + errPacket);
        }

        ResultSetHeaderPacket rsHeader = new ResultSetHeaderPacket();
        rsHeader.fromBytes(body);

        List<FieldPacket> fields = new ArrayList<>();
        for (int i = 0; i < rsHeader.getColumnCount(); i++) {
            FieldPacket fp = new FieldPacket();
            fp.fromBytes(readNextPacket());
            fields.add(fp);
        }

        readEofPacket();

        List<RowDataPacket> rowData = new ArrayList<>();
        while (true) {
            body = readNextPacket();
            if (body[0] == -2) {
                break;
            }
            RowDataPacket rowDataPacket = new RowDataPacket();
            rowDataPacket.fromBytes(body);
            rowData.add(rowDataPacket);
        }

        ResultSetPacket resultSet = new ResultSetPacket();
        resultSet.getFieldDescriptors().addAll(fields);
        for (RowDataPacket r : rowData) {
            resultSet.getFieldValues().addAll(r.getColumns());
        }
        resultSet.setSourceAddress(channel.getRemoteSocketAddress());

        return resultSet;
    }

    private boolean readEofPacket() throws IOException {
        byte[] eofBody = readNextPacket();
        EOFPacket packet = new EOFPacket();
        packet.fromBytes(eofBody);
        if (eofBody[0] != -2) {
            throw new IOException("EOF Packet is expected, but packet with field_count=" + eofBody[0] + " is found.");
        }

        return (packet.getStatusFlag() & 0x0008) != 0;
    }
}